//
// Created on 2025/6/15.
//
// Node APIs are not fully supported. To solve the compilation error of the interface cannot be found,
// please include "napi/native_api.h".

#ifndef VARINT_DECODE_H
#define VARINT_DECODE_H

#include "common.h"
#include "varint/stream.h"
#include <algorithm>
#include <stdexcept>


inline uint16_t extract_msb(uint8x16_t vec) {
    // 将每个元素的最高位移至最低位，结果0或1
    uint8x16_t shifted = vshrq_n_u8(vec, 7);

    // 分割为低8位和高8位
    uint8x8_t low = vget_low_u8(shifted);
    uint8x8_t high = vget_high_u8(shifted);

    // 权值数组，对应每个字节的位位置
    const uint8_t weights[8] = {1, 2, 4, 8, 16, 32, 64, 128};
    uint8x8_t weights_vec = vld1_u8(weights);

    // 计算低8位和高8位的掩码
    uint16_t low_mask = vaddlv_u8(vmul_u8(low, weights_vec));
    uint16_t high_mask = vaddlv_u8(vmul_u8(high, weights_vec));

    // 合并成最终的16位掩码
    return low_mask | (high_mask << 8);
}

inline void clear_msb(uint8x16_t value, uint8_t out[16]) { vst1q_u8(out, vandq_u8(value, vdupq_n_u8(0x7f))); }

inline uint64_t decode_varint(const uint8_t *buf) {
    uint64_t result = 0;
    do {
        result |= (*buf & 0x7f);
        result <<= 7;
    } while (*buf++ & 0x80);
    return result;
}

// 传入已清除标志位缓冲
inline uint64_t decode_varint(const uint8_t *buf, uint8_t len) {
    uint64_t result = 0;
    switch (len) {
    case 10:
        result |= (uint64_t)buf[9] << 63;
        [[fallthrough]];
    case 9:
        result |= (uint64_t)buf[8] << 56;
        [[fallthrough]];
    case 8:
        result |= (uint64_t)buf[7] << 49;
        [[fallthrough]];
    case 7:
        result |= (uint64_t)buf[6] << 42;
        [[fallthrough]];
    case 6:
        result |= (uint64_t)buf[5] << 35;
        [[fallthrough]];
    case 5:
        result |= (uint64_t)buf[4] << 28;
        [[fallthrough]];
    case 4:
        result |= (uint64_t)buf[3] << 21;
        [[fallthrough]];
    case 3:
        result |= (uint64_t)buf[2] << 14;
        [[fallthrough]];
    case 2:
        result |= (uint64_t)buf[1] << 7;
        [[fallthrough]];
    case 1:
        result |= (uint64_t)buf[0];
        [[fallthrough]];
    default:
        return result;
    }
}


inline uint64_t decode_varint(const uint8_t **buf) {
    uint64_t result = 0;
    uint8_t shift = 0;
    do {
        result |= (**buf & 0x7f) << shift;
        shift += 7;
    } while (**buf++ & 0x80);
    return result;
}

template <typename T> static void varint_decode(const uint8_t *input, size_t input_len, T *output) {
    const uint8_t *end = input + input_len;
    using unsigned_type = typename std::make_unsigned<T>::type;

    alignas(16) uint8_t temp[16];

    while (input + 16 < end) {
        uint8x16_t value = vld1q_u8(input);
        uint16_t mask = extract_msb(value); // 获取掩码图
        uint16_t inv_mask = ~mask;
        if (inv_mask == 0) // varint不可能存在16字节没有停止符
            return;
        clear_msb(value, temp); // simd清理高位标志位
        int count = __builtin_popcount(inv_mask);
        uint8_t shift = 0;
        for (int i = 0; i < count; i++) {
            uint8_t n = __builtin_ctz(inv_mask) + 1;
            unsigned_type val = (unsigned_type)decode_varint(temp + shift, n);
            if (std::is_same_v<T, int64_t> || std::is_same_v<T, int32_t>) {
                *output++ = zigzag_decode(val);
            } else {
                *output++ = val;
            }
            shift += n;
            inv_mask >>= n;
        }
        input += shift;
    }

    while (input < end) {
        uint32_t result = 0;
        uint8_t shift = 0;
        do {
            result |= ((T)*input & 0x7f) << shift;
            shift += 7;
        } while (*input++ & 0x80);
        *output++ = result;
    }
}


/**
 * varint uint32解码
 * @param stream
 * @return
 */
static OP_RESULT varint_decode(varint_stream *stream) {
    const uint8_t *end = stream->next_in + stream->avail_in;
    alignas(16) uint8_t temp[16];

    if (stream->avail_out < sizeof(uint32_t))
        return OP_RESULT::OUTPUT_TOO_SMALL;

    if (stream->pos > 0) {
        bool success = false;
        while (stream->avail_in) {
            if (stream->pos > 5)
                return OP_RESULT::BAD_DATA;

            uint8_t val = *stream->next_in;
            stream->value |= ((val & 0x7f) << (stream->pos * 7));
            stream->next_in++;
            stream->avail_in--;
            stream->pos++;
            if (!(val & 0x80)) {
                *(uint32_t *)stream->next_out = stream->value;
                stream->next_out += sizeof(uint32_t);
                stream->avail_out -= sizeof(uint32_t);
                stream->pos = 0;
                success = true;
                break;
            }
        }
        if (!success)
            return OP_RESULT::NEED_MORE_INPUT;
    }

    // simd + 掩码批量处理16字节
    while ((stream->next_in + 16) < end) {
        uint8x16_t value = vld1q_u8(stream->next_in);

        // 获取掩码图
        uint16_t mask = extract_msb(value);
        uint16_t inv_mask = ~mask;

        // varint不可能存在16字节没有停止符
        if (inv_mask == 0)
            return OP_RESULT::BAD_DATA;

        // simd清理高位标志位
        clear_msb(value, temp);
        int varint_count = __builtin_popcount(inv_mask);

        // 预防输出缓冲区不足, uint32_t和uint64_t是2的倍数，底层应该会进行右移优化
        int count = std::min<int>(varint_count, stream->avail_out / sizeof(uint32_t));
        uint8_t shift = 0;
        for (int i = 0; i < count; i++) {
            uint8_t n = __builtin_ctz(inv_mask) + 1;
            uint32_t val = decode_varint(temp + shift, n);
            memcpy(stream->next_out, &val, sizeof(uint32_t));
            shift += n;
            inv_mask >>= n;
        }
        stream->next_in += shift;
        stream->avail_in -= shift;
        stream->next_out += sizeof(uint32_t) * count;
        stream->avail_out += sizeof(uint32_t) * count;
        if (varint_count != count)
            return OP_RESULT::OUTPUT_TOO_SMALL;
    }

    memset(temp, 0xff, 16);
    memcpy(temp, stream->next_in, stream->avail_in);
    uint8x16_t value = vld1q_u8(temp);
    // 获取掩码图
    uint16_t mask = extract_msb(value);
    clear_msb(value, temp);

    uint16_t inv_mask = ~mask;
    int varint_count = __builtin_popcount(inv_mask);
    int count = std::min<int>(varint_count, stream->avail_in / sizeof(uint32_t));

    uint32_t p = 0;
    for (int i = 0; i < count; i++) {
        uint8_t n = __builtin_ctz(inv_mask) + 1;
        uint32_t val = decode_varint(temp + p, n);
        memcpy(stream->next_out, &val, sizeof(uint32_t));
        p += n;
        inv_mask >>= n;
    }

    stream->next_in += p;
    stream->avail_in -= p;
    stream->next_out += sizeof(uint32_t) * count;
    stream->avail_out += sizeof(uint32_t) * count;

    if (varint_count != count) {
        return OP_RESULT::OUTPUT_TOO_SMALL;
    }

    if (stream->next_in < end) {
        stream->pos = 0;
        stream->value = 0;
        while (stream->next_in < end) {
            stream->value |= (uint32_t)((*stream->next_in) & 0x7f) << (stream->pos * 7);
            stream->pos++;
        }
        return OP_RESULT::NEED_MORE_INPUT;
    }


    return OP_RESULT::SUCCESS;
}


#endif // VARINT_DECODE_H
